from datetime import datetime,timedelta
import pandas as pd

class Operation:
    def sell(self,context, invest_code, invest_data, dt):
        """ 平仓卖出投资标的 """
        sell_price = invest_data[invest_code].loc[dt]
        context.currency += context.now_have[invest_code]['position'] * sell_price
        context.asset -= context.now_have[invest_code]['position'] * sell_price
        del context.now_have[invest_code]
        self.update_account(context,invest_data,dt)

    def buy(self,context, invest_code, money_per_invest, invest_data, dt):
        """ 开仓买入投资标的 """
        buy_price = invest_data[invest_code].loc[dt]
        amount = money_per_invest / buy_price
        if amount>100: # 100份以上才能开仓，小于100份会开仓失败
            context.currency -= money_per_invest
            context.asset += money_per_invest
            if invest_code not in context.now_have:
                context.now_have[invest_code] = {}
                context.now_have[invest_code]['position'] = 0
            context.now_have[invest_code]['position'] += amount
            context.now_have[invest_code]['dt'] = dt
            context.now_have[invest_code]['price'] = buy_price
            self.update_account(context, invest_data, dt)

    def adjust_have(self,context,dt,invest_data):
        """ 调整仓位 """
        if list(context.now_have.keys())==context.target_have:
            # 如果今天的目标仓位和已有仓位相同，则不做任何操作
            return 1
        else:
            # 先卖后买，注意两个list为空的情况
            for i_now in list(context.now_have.keys()):
                if i_now not in context.target_have:
                    self.sell(context,i_now,invest_data,dt)
            if len(context.target_have)>0:
                money_per_invest = context.currency / len(set(context.target_have)-set(context.now_have.keys()))
                for j_target in context.target_have:
                    if j_target not in context.now_have:
                        self.buy(context,j_target,money_per_invest,invest_data,dt)

    def update_account(self,context,invest_data,dt):
        """ 每一天更新账户的钱 """
        if len(context.now_have)>0:
            temp_asset = 0
            for i_now in context.now_have:
                i_price = invest_data[i_now].loc[dt]
                temp_asset += i_price * context.now_have[i_now]['position']
            context.asset = temp_asset
            context.total = context.currency + context.asset